package geppetto.cat.programs;


import geppetto.cat.common.StaticTools;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.AgreementM1;
import geppetto.cat.models.HMM;
import geppetto.cat.models.M1;
import geppetto.cat.models.SubstochasticHMM;
import geppetto.cat.models.SubstochasticM1;

import java.io.IOException;



/** class to train and save an (HMM) alignment model */
public class SaveModel {

	public static void main(String[] args) throws IOException {
		String corpusFile = args[0];
		int size = Integer.parseInt(args[1]); // 100k
		int maxSentenceSize = Integer.parseInt(args[2]); // 40
		int numberIterations = Integer.parseInt(args[3]); // 5
		String dir = args[4];
		String modelName = args[5];
		System.out.println("Saving Models experiment: ");
		System.out.println("Corpus " + corpusFile);
		System.out.println("Size " + size);
		System.out.println("Max Sentence size " + maxSentenceSize);
		System.out.println("Number of iterations " + numberIterations);
		System.out.println("OutputDir " + dir);
		System.out.println("Model Type " + modelName);
		System.out.println("---");
		BilingualCorpus corpusF = BilingualCorpus.getCorpusFromFileDescription(
				corpusFile, size, maxSentenceSize);
		BilingualCorpus corpusB = corpusF.reverse();
		String baseDir = dir + modelName + "/" + corpusF.getName() + "/" + size;
		String modelDir = baseDir + "/model/";

		if (modelName.equalsIgnoreCase("baseline")) {
			M1 m1F;
			M1 m1B;
			HMM mhmmF;
			HMM mhmmB;

			m1F = new M1(corpusF);
			m1F.train(numberIterations,false,"");
			m1B = new M1(corpusB);
			m1B.train(numberIterations,false,"");

			String M1Dir = modelDir + "/M1/";
			StaticTools.createDir(M1Dir);
			m1F.saveModel(M1Dir + "forward");
			m1B.saveModel(M1Dir + "backward");

			// NOTE: we're intentionally clobbering the translation table of
			// model1! (since we aren't using it anymore)
			mhmmF = new HMM(corpusF, m1F._tb);
			mhmmF.train(numberIterations,false,"");
			mhmmB = new HMM(corpusB, m1B._tb);
			mhmmB.train(numberIterations,false,"");

			String MHMMDir = modelDir + "/MHMM/";
			StaticTools.createDir(MHMMDir);
			mhmmF.saveModel(MHMMDir + "forward");
			mhmmB.saveModel(MHMMDir + "backward");
			System.out.println("SAVE MODEL Experience");

		} else if (modelName.equalsIgnoreCase("agreement")) {

			int projectonIterations = 5;
			double epsilon = 0.0;
			AgreementM1 m1;
			AgreementHMM mhmm;
			m1 = new AgreementM1(corpusF, corpusB, epsilon, projectonIterations);
			m1.train(numberIterations);
			String M1Dir = modelDir + "/M1/";
			StaticTools.createDir(M1Dir);
			m1.saveModel(M1Dir + "forward");
			
			mhmm = new AgreementHMM(corpusF, corpusB, m1.forward._tb,m1.backward._tb, epsilon, projectonIterations);
			mhmm.train(numberIterations);
			String MHMMDir = modelDir + "MHMM/";
			StaticTools.createDir(MHMMDir);
			mhmm.saveModel(MHMMDir);
			System.out.println("Done with the SAVE MODEL Experience");
		} else if (modelName.equalsIgnoreCase("substochastic")) {
			SubstochasticM1 m1F;
			SubstochasticM1 m1B;
			SubstochasticHMM mhmmF;
			SubstochasticHMM mhmmB;

			m1F = new SubstochasticM1(corpusF,5);
			m1F.train(numberIterations,false,"");
			m1B = new SubstochasticM1(corpusB,5);
			m1B.train(numberIterations,false,"");

			String M1Dir = modelDir + "/M1/";
			StaticTools.createDir(M1Dir);
			m1F.saveModel(M1Dir + "forward");
			m1B.saveModel(M1Dir + "backward");

			// NOTE: we're intentionally clobbering the translation table of
			// model1! (since we aren't using it anymore)
			mhmmF = new SubstochasticHMM(corpusF, m1F._tb,5);
			mhmmF.train(numberIterations,false,"");
			mhmmB = new SubstochasticHMM(corpusB, m1B._tb,5);
			mhmmB.train(numberIterations,false,"");

			String MHMMDir = modelDir + "/MHMM/";
			StaticTools.createDir(MHMMDir);
			mhmmF.saveModel(MHMMDir + "forward");
			mhmmB.saveModel(MHMMDir + "backward");
			System.out.println("Done with the SAVE MODEL Experience");
		} else {
			System.out.println("Unknown Model" + modelName);
			System.exit(1);
		}
	}

}
